from sentence_transformers import SentenceTransformer
from pymilvus import connections, Collection, CollectionSchema, FieldSchema, DataType, utility
import numpy as np
import os

class SentenceMilvusManager:
    def __init__(self, model_name='all-MiniLM-L6-v2', collection_name='sentence_collection'):
        """
        初始化Sentence-BERT和Milvus Lite管理器
        
        Args:
            model_name: Sentence-BERT模型名称
            collection_name: Milvus集合名称
        """
        print("正在初始化模型和数据库连接...")
        self.model = SentenceTransformer(model_name)
        self.collection_name = collection_name
        self.vector_dim = 384  # all-MiniLM-L6-v2模型的向量维度
        
        # 确保milvus_data目录存在
        milvus_path = os.path.join(os.getcwd(), "milvus_data")
        os.makedirs(milvus_path, exist_ok=True)
        
        # 设置数据库文件路径
        db_path = os.path.join(milvus_path, "milvus.db")
        
        # 连接到Milvus Lite
        connections.connect(
            alias="default",
            uri=db_path,  # 使用具体的.db文件路径
            local_path=milvus_path
        )
        
        # 如果集合已存在，先删除它
        if utility.has_collection(self.collection_name):
            utility.drop_collection(self.collection_name)
        
        # 创建或获取集合
        self._init_collection()
    
    def _init_collection(self):
        """初始化Milvus集合"""
        if utility.has_collection(self.collection_name):
            self.collection = Collection(self.collection_name)
            print(f"已连接到现有集合: {self.collection_name}")
            return

        # 定义集合字段
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=500),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.vector_dim)
        ]
        
        # 创建集合schema
        schema = CollectionSchema(fields=fields, description="句子向量集合")
        
        # 创建集合
        self.collection = Collection(name=self.collection_name, schema=schema)
        
        # 创建索引
        index_params = {
            "metric_type": "COSINE",  # 使用余弦相似度
            "index_type": "HNSW",     # 使用HNSW索引
            "params": {
                "M": 8,               # HNSW图的最大度数
                "efConstruction": 64   # 构建时的搜索深度
            }
        }
        self.collection.create_index(
            field_name="embedding",
            index_params=index_params,
            index_name="embedding_index"
        )
        print(f"已创建新集合: {self.collection_name}")
        
        # 加载集合
        self.collection.load()
    
    def insert_sentences(self, sentences):
        """
        将句子转换为向量并插入Milvus
        
        Args:
            sentences: 句子列表
        """
        print("正在处理句子...")
        # 编码句子
        embeddings = self.model.encode(sentences)
        
        # 确保向量数据是浮点数格式
        embeddings = embeddings.astype('float32')
        
        # 准备插入数据
        data = [
            {
                'text': text,
                'embedding': embedding.tolist()
            }
            for text, embedding in zip(sentences, embeddings)
        ]
        
        # 插入数据
        self.collection.insert(data)
        self.collection.flush()
        print(f"已插入 {len(sentences)} 条数据")
    
    def search_similar_sentences(self, query_sentence, top_k=3):
        """
        搜索与给定句子最相似的句子
        
        Args:
            query_sentence: 查询句子
            top_k: 返回最相似的前k个结果
            
        Returns:
            list: 包含相似句子和距离的列表
        """
        print(f"正在搜索与'{query_sentence}'最相似的{top_k}个句子...")
        
        # 加载集合
        self.collection.load()
        
        # 编码查询句子
        query_embedding = self.model.encode([query_sentence])
        
        # 确保向量数据是浮点数格式
        query_embedding = query_embedding.astype('float32')
        
        # 搜索
        search_params = {
            "metric_type": "COSINE",  # 使用余弦相似度
            "params": {
                "ef": 32  # 搜索时的候选邻居数
            }
        }
        
        results = self.collection.search(
            data=[query_embedding[0].tolist()],  # 只取第一个向量，并转换为列表
            anns_field="embedding",
            param=search_params,
            limit=top_k,
            output_fields=["text"]
        )
        
        # 释放集合
        self.collection.release()
        
        # 整理结果
        similar_sentences = []
        for hits in results:
            for hit in hits:
                similar_sentences.append({
                    'text': hit.entity.get('text'),
                    'distance': hit.distance
                })
        
        return similar_sentences

def main():
    # 创建管理器实例
    manager = SentenceMilvusManager()
    
    # 示例句子
    sentences = [
        "我喜欢在公园散步",
        "公园里散步很舒服",
        "今天天气真不错",
        "这个苹果很甜",
        "我最喜欢吃水果了",
        "公园里的风景很美",
        "散步对健康有好处",
        "水果对身体有益",
        "今天阳光明媚",
        "我喜欢户外运动"
    ]
    
    # 插入句子
    manager.insert_sentences(sentences)
    
    # 测试相似度搜索
    query_sentences = [
        "公园散步",
        "新鲜水果",
        "户外活动"
    ]
    
    # 对每个查询句子进行相似度搜索
    for query in query_sentences:
        print("\n" + "="*50)
        similar = manager.search_similar_sentences(query, top_k=3)
        print(f"\n与'{query}'最相似的句子：")
        for idx, item in enumerate(similar, 1):
            print(f"{idx}. {item['text']} (距离: {item['distance']:.4f})")

if __name__ == "__main__":
    main()
